"""
文本主题分类: cardiffnlp/tweet-topic-latest-multi
共19个类别
"""
import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer
)
from scipy.special import expit
from threading import RLock
import numpy as np


class TextClassifier(object):

    lock = RLock()
    def __new__(cls, *args, **kwargs):
        with TextClassifier.lock:
            if not hasattr(TextClassifier, "_instance"):
                TextClassifier._instance = object.__new__(cls)
                repo_id = "cardiffnlp/tweet-topic-latest-multi"
                cls.tokenizer = AutoTokenizer.from_pretrained(repo_id)
                model = AutoModelForSequenceClassification.from_pretrained(repo_id)
                cls.model = model
                cls.class_mapping = model.config.id2label
        return TextClassifier._instance

    def predict(self, text):
        tokens = self.tokenizer(text[0:512], return_tensors='pt')
        output = self.model(**tokens)

        scores = output[0][0].detach().numpy()
        scores = expit(scores)
        # predictions = (scores >= 0.5) * 1
        indexes = np.asarray(np.where(scores >= 0.5)).reshape(-1)
        return [self.class_mapping.get(idx) for idx in indexes]


if __name__ == '__main__':
    classifier = TextClassifier()
    text = "It is great to see athletes promoting awareness for climate change."
    class_names = classifier.predict(text)
    print(class_names)
